import sqlite3
from pathlib import Path
import sys
import argparse
import zipfile
import os
import hashlib

def check_and_connect_db(db_path):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # 设置返回字典格式的结果
    cursor = conn.cursor()
    return conn, cursor

def calculate_file_hash(file_content):
    """计算文件内容的SHA-256哈希值"""
    sha256_hash = hashlib.sha256()
    sha256_hash.update(file_content)
    return sha256_hash.hexdigest()

def process_zip_file(zip_path, file_name, file_hash_to_find, cursor, conn, record_id):
    with zipfile.ZipFile(zip_path, 'r') as zipf:
        # 首先尝试根据文件名查找
        if file_name in zipf.namelist():
            file_content = zipf.read(file_name)
            file_hash = calculate_file_hash(file_content)
            if file_hash == file_hash_to_find:
                cursor.execute("UPDATE photo_move SET deal_flag=9 WHERE id=?", (record_id,))
                conn.commit()
                print(f"已更新deal_flag为9: 记录ID {record_id}, 找到匹配的file_hash: {file_hash_to_find} 对应文件名 {file_name}")
                return True
        
        # 如果根据文件名没有找到匹配项，则遍历所有文件查找相同的file_hash
        for file_info in zipf.filelist:
            file_content = zipf.read(file_info.filename)
            file_hash = calculate_file_hash(file_content)
            if file_hash == file_hash_to_find:
                cursor.execute("UPDATE photo_move SET deal_flag=9 WHERE id=?", (record_id,))
                conn.commit()
                print(f"已更新deal_flag为9: 记录ID {record_id}, 找到匹配的file_hash: {file_hash_to_find} 在文件 {file_info.filename}")
                return True
        
        print(f"未找到匹配的file_hash: {file_hash_to_find} 在ZIP文件 {zip_path}")
        return False

def process_records(db_directory, db_conn, cursor):
    print("开始处理 deal_flag=4 的记录")
    
    # 查询photo_move表中所有deal_flag为4的记录
    cursor.execute("SELECT * FROM photo_move WHERE deal_flag=4 ORDER BY id ASC")  # 假设id字段用于排序
    records = cursor.fetchall()

    if not records:
        print("没有需要处理的记录")
        return
    
    for record in records:
        aim_zip = record['aim_zip']
        zip_path = Path(db_directory) / aim_zip
        file_hash_to_find = record['file_hash']
        file_name = record['file_name']  # 获取文件名
        
        if zip_path.exists() and zipfile.is_zipfile(zip_path):
            process_zip_file(zip_path, file_name, file_hash_to_find, cursor, db_conn, record['id'])
        else:
            print(f"警告: ZIP文件 {zip_path} 不存在或不是有效的ZIP文件")

def main(db_directory):
    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn, cursor = check_and_connect_db(db_path)

    process_records(db_directory, conn, cursor)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Check records with deal_flag=4, verify file_hash within the corresponding ZIP files.")
    parser.add_argument('--db', required=True, help="Directory containing the SQLite database file (file_hashes.db)")

    args = parser.parse_args()

    db_directory = args.db

    if not os.path.isdir(db_directory):
        print(f"错误: {db_directory} 不是一个有效的目录")
        sys.exit(1)

    main(db_directory)